import argparse
import logging
from collections import namedtuple

import torch
from colorlog import ColoredFormatter
import os

args = argparse.ArgumentParser()
args.add_argument('-dataset', default='power',
                  help="used dataset")
args.add_argument('-random_seed', type=int, default=666)

# model related
args.add_argument('-infect_rate', type=float, default=0.15)
args.add_argument('-recover_rate', type=float, default=0.04, help="used in SIR model")
args.add_argument('-generate_step', type=int, default=60, help="iteration step of label data generation")
args.add_argument('-generate_model', type=str, default="SI", help="one of SI, SIR and IC")
args.add_argument('-num_src', type=int, default=500)
args.add_argument('-num_src_ratio', type=float, default=0.01)
args.add_argument('-lb', type=float, default=0.2, help="lower bound of of the infection ratio")
args.add_argument('-ub', type=float, default=0.7, help="upper bound of of the infection ratio")
args.add_argument('-seq_len', type=int, default=30)

args.add_argument('-lpsi_alpha', type=float, default=0.2)
args.add_argument('-cluster_number', type=int, default=1)
args.add_argument('-lpsi_use_iter', type=bool, default=True)
args.add_argument('-lpsi_iter_step', type=int, default=7)

args.add_argument('-gcn_layers', type=int, default=3)
args.add_argument('-gcn_hidden', type=int, default=128)

args.add_argument("-eval_step", type=int, default=10)

args.add_argument("-acc_tol", type=int, default=1)
args.add_argument("-eliminate_rate", type=float, default=0)
args.add_argument("-eval_data_ratio", type=float, default=0.1)
args.add_argument("-data_group_size", type=int, default=100)
args.add_argument('-batch_size', type=int, default=5)
args.add_argument("-calc_erd", type=bool, default=True)
args.add_argument("-max_erd", type=int, default=10)

# transformer related
args.add_argument("-d_model", type=int, default=512)
args.add_argument("-nhead", type=int, default=8)
args.add_argument("-decoder_layers", type=int, default=3)

# 目前先不使用动态图，先focus on静态图partial observation
args.add_argument("-use_dynamic_graph", type=bool, default=True, help="treat dynamic graph as dynamic graph")
args.add_argument("-dgraph_start", type=float, default=0.1, help="ratio of time that has passed in dynamic graph "
                                                                 "when the propagation begins")
args.add_argument("-lp_duration", type=float, default=0.9, help="ratio of propagation period and all dgraph period")

# load & restore option
args.add_argument('-use_cache', type=bool, default=True)
args.add_argument('-update_cache', type=bool, default=False)

# device option
args.add_argument('-use_cuda', type=bool, default=True)

# train option
args.add_argument('-learning_rate', type=float, default=0.001)
args.add_argument('-epochs', type=int, default=100)
args.add_argument('-kfolds', type=int, default=10, help="number of folds")
args.add_argument('-dropout', type=float, default=0.1)
args.add_argument('-weight_decay', type=float, default=1e-8)



args = args.parse_args()

# 创建一个日志格式器，其中包含颜色设置
formatter = ColoredFormatter(
    "[%(levelname)s] %(asctime)s %(message)s",
    datefmt=None,
    reset=True,
    log_colors={
        'DEBUG': 'cyan',
        'INFO': 'green',
        'WARNING': 'yellow',
        'ERROR': 'red',
        'CRITICAL': 'red,bg_white',
    },
    secondary_log_colors={},
    style='%'
)

# 配置基本日志设置
handler = logging.StreamHandler()
handler.setFormatter(formatter)
logging.basicConfig(level=logging.INFO, handlers=[handler])

DeviceDict = namedtuple('DeviceDict', ['device', 'backup_device'])

if torch.cuda.is_available() and args.use_cuda:
    device = 'cuda'
elif torch.backends.mps.is_available():
    device = 'mps'
else:
    device = 'cpu'
backup_device = 'cpu'
devs = DeviceDict(device=device, backup_device=backup_device)

cur_cwd = os.getcwd()
suffix_to_change = ['/src', '\\src']
for suffix in suffix_to_change:
    if cur_cwd.endswith(suffix):
        os.chdir(cur_cwd.removesuffix(suffix))

if __name__ == '__main__':
    print(args)

    for k, v in args.__dict__.items():
        logging.info(f'__{k}: {v}')
